import os
from typing import Annotated, Any, Dict, Generic, Optional, TypeVar

from mcp.server.fastmcp import FastMCP
from pydantic import BaseModel, Field

from lc_radar_mcp.common.config import load_config
from lc_radar_mcp.common.logger import get_logger
from lc_radar_mcp.services.lc_radar import WeatherPerceptionSDK

# 优先从环境变量读取API_KEY，如果没有则从配置文件读取
env_api_key = os.getenv("API_KEY")
config = load_config("config.yaml")
if env_api_key:
    config["api_key"] = env_api_key

# 初始化 FastMCP 服务
mcp = FastMCP("weather-perception", description="Weather Perception MCP 服务", version="1.0.0")
sdk = WeatherPerceptionSDK(config=config, logger=get_logger(name="wp_sdk"))
logger = get_logger(name="weather-perception")

T = TypeVar("T")

class ApiResponse(BaseModel, Generic[T]):
    success: bool
    data: Optional[T] = None
    error: Optional[str] = None
    meta: Optional[Dict[str, Any]] = None

    @classmethod
    def ok(cls, data: T, meta: Dict[str, Any] = None) -> "ApiResponse[T]":
        return cls(success=True, data=data, meta=meta)

    @classmethod
    def fail(cls, error: str, meta: Dict[str, Any] = None) -> "ApiResponse[None]":
        return cls(success=False, error=error, meta=meta)

# MCP Prompt
@mcp.prompt(name="assistant", description="气象站智能查询助手，支持站点数据检索与多参数筛选")
def weather_assistant(query: str) -> str:
    return (
        "你是气象站智能数据助手，支持指定站点编号快速检索观测数据。"
        "当用户提供站点编号时，调用 query_station 工具检索该站点观测信息。"
        "可以进一步用 page_num、page_size 控制返回条数。"
        f"用户需求：{query}"
    )

@mcp.tool(name="query_station", description="根据站点编号获取 weather-perception 实时/历史观测数据")
async def query_station(
    devicename: Annotated[str, Field(description="站点编号，如 'Z9531'")],
    page_num: Annotated[int, Field(description="页码，从1开始", ge=1)] = 1,
    page_size: Annotated[int, Field(description="每页数量，默认240")] = 240,
) -> ApiResponse:
    """
    查询站点观测数据

    Args:
        devicename (str): 站点编号
        page_num (int): 页码，默认1
        page_size (int): 每页条数，默认240
    Returns:
        ApiResponse: 查询结果
    """
    logger.info(f"Querying weather-perception: devicename={devicename}, page_num={page_num}, page_size={page_size}")
    try:
        result = await sdk.query_station(devicename=devicename, page_num=page_num, page_size=page_size)
        if not result:
            return ApiResponse.fail("查询结果为空，请检查日志。")
        logger.info(f"query_station result: {result}")
        return ApiResponse.ok(data=result, meta={
            "devicename": devicename,
            "page_num": page_num,
            "page_size": page_size
        })
    except Exception as e:
        logger.error(f"Error querying station {devicename}: {e}")
        return ApiResponse.fail(str(e))
